import json
from itertools import chain

from ansi2html import Ansi2HTMLConverter
from loguru import logger
from nicegui import ui
from rr.api import ServerHelper, TaskStatus

server_helper = ServerHelper()
conv = Ansi2HTMLConverter(
    line_wrap=True,
    inline=False,
    dark_bg=True,
    escaped=False,
)


field_renderers = {}
live_updates_session = {}


def render_graph_for_task(task_hash: str):
    """
    :param task_hash:
    :return:
    """
    """
    digraph UML_Class_diagram {
    node [
        shape=record
        style=filled
        fillcolor=gray95
    ]



    System_1 [
        shape=plain
        label=<<table border="1" cellborder="0" cellspacing="0" cellpadding="1">
            <tr> <td> <b>System</b> </td> </tr>
            <tr> <td align="left" >+ property</td> </tr>
            <tr> <td port="ss1" align="left" >- p1</td> </tr>
            <tr> <td port="ss2" align="left" >- p2</td> </tr>
            <tr> <td port="ss3" align="left" >- p3</td> </tr>
            <tr> <td align="left">...</td> </tr>
        </table>>
    ]

    System_1:ss2 -> Subsystem_2;
    System_1:ss3 -> Subsystem_3;

    }
    """
    pass


def register_renderer(field_name: str):
    def wrapper(cls):
        field_renderers[field_name] = cls
        return cls

    return wrapper


@register_renderer("text/html")
def render_text_html(text: str):
    return ui.html(text)


@register_renderer("image/png")
def render_png(text: str):
    # with ui.element(tag="div"):
    ui.image("data:image/png;base64," + text).classes("min-width-96 max-w-xl")


@register_renderer("text/plain")
def render_plain(text: str):
    with ui.element(tag="pre") as ele:
        ui.html(text)
        return ele


@register_renderer("text/json")
def render_json(text: str):
    with ui.json_editor({"content": {"json": json.loads(text)}}) as ele:
        return ele


@register_renderer("text/mermaid")
def render_mermaid(text: str):
    # ui.code(text, language="mermaid").classes("my-5")
    return ui.mermaid(text)


@register_renderer("text/markdown")
def render_markdown(text: str):
    return ui.markdown(text)


@register_renderer("stdout")
def render_stdout(text: str):
    return ui.html(conv.convert(text.strip()))


@register_renderer("stderr")
def render_stderr(text: str):
    return ui.html(conv.convert(text.strip()))


def ui_task_detail(task_id):
    assert task_id is not None
    assert isinstance(task_id, str)
    logger.info(f"ui_task_detail {task_id}")
    _ui_task_detail(task_id)


def render_single_element(item):
    if item["output_type"] == "display_data":
        for mime_type, v in item["data"].items():
            # logger.info(f"mime_type: {mime_type}")
            field_renderers[mime_type](v)
    elif item["output_type"] == "stream":
        render_stdout(item["text"])
    elif item["output_type"] == "error":
        render_stdout("\n".join(item["traceback"]))


def get_output_iter(task_id, limit: int = 100, offset: int = 0):
    logs = server_helper.get_logs(task_id, limit, offset)


@ui.refreshable
def _ui_task_detail(task_id, element=None):
    call_dict = dict(
        offset=0,
        limit=10,
    )

    def runrun():
        try:
            cnt_this_time = 0
            for record in server_helper.get_logs(task_id, **call_dict):
                render_single_element(record)
                cnt_this_time += 1
            call_dict["offset"] += cnt_this_time
            if (
                cnt_this_time == 0
                and server_helper.get_task_status(task_id) != TaskStatus.RUNNING
            ):
                logger.info("task finished")
                timer.cancel()
        except StopIteration:
            timer.cancel()

    timer = ui.timer(0.5, runrun)

    # Add files section
    with ui.card().classes("w-full"):
        with ui.expansion("Task Files", icon="folder").classes(
            "text-xl font-bold w-full"
        ):
            files = server_helper.get_task_files(task_id)
            if not files:
                ui.label("No files found")
            else:

                def on_click(event):
                    file_id = event.args["data"]["file_id"]
                    column_id = event.args["colId"]
                    if column_id == "download":
                        ui.navigate.to(f"/file/{file_id}")
                    elif column_id == "copy_link":
                        link = f"/file/{file_id}"
                        ui.clipboard.write(link)
                        ui.notify(f"File viewer link copied to clipboard")
                    elif column_id == "file_id":
                        ui.navigate.to(f"/file/{file_id}")

                row_data = [file.model_dump() for file in files]
                for item in row_data:
                    item["download"] = "view"
                    item["copy_link"] = "link"

                ui.aggrid(
                    {
                        "columnDefs": [
                            {
                                "headerName": "File ID",
                                "field": "file_id",
                                "sortable": True,
                            },
                            {"headerName": "Type", "field": "type_", "sortable": True},
                            {
                                "headerName": "Content Type",
                                "field": "content_type",
                                "sortable": True,
                            },
                            {
                                "headerName": "View",
                                "field": "download",
                                "sortable": False,
                                "width": 30,
                            },
                            {
                                "headerName": "Copy Link",
                                "field": "copy_link",
                                "sortable": False,
                                "width": 30,
                            },
                        ],
                        "rowData": row_data,
                        "domLayout": "autoHeight",
                    },
                    html_columns=[3, 4],  # Make view and copy link columns clickable
                ).on("cellClicked", on_click)


page_config = {
    "nav_entry": [
        {
            "title": None,
            "url": "/task_detail/{task_id}",
            "function": ui_task_detail,
        }
    ]
}
